/*
** Copyright (C) 2009 Zhang Cong <ftofficer@ftofficer.com>
**  
** This program is free software; you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation; either version 2 of the License, or
** (at your option) any later version.
**  
** This program is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**  
** You should have received a copy of the GNU General Public License
** along with this program; if not, write to the Free Software
** Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
**  
*/

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif // HAVE_CONFIG_H

#include "portable.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

#include <libssh2.h>
#include <sys/types.h>
#include <sys/stat.h>

#include "liboi.h"
#include "liboi_log.h"
#include "oissh.h"

typedef struct oi_ssh_device_data
{
    socket_t ssh_sock;
    LIBSSH2_SESSION* ssh_session;

    oi_device_info dev_info;

} oi_ssh_device_data;

static int oi_ssh_report_lasterror(const char* file, int line,
                                    const char* func_name,
                                    oi_ssh_device_data* dev_data,
                                    const char* what)
{
    char* errmsg = NULL;
    int errmsg_len = 0;
    int err_no = 0;

    err_no = libssh2_session_last_error(dev_data->ssh_session,
                                        &errmsg, &errmsg_len, 0);
    oi_log_write(file, line, LOG_ERROR, 
                 "%s: %s: %d, %s",
                 func_name, what, err_no, errmsg);
    return err_no;
}

#define OI_SSH_LASTERROR(func_name, dev_data, what) \
    oi_ssh_report_lasterror(__FILE__, __LINE__, func_name, dev_data, what)

static oi_ssh_device_data* create_device_data()
{
    const char* my_name = "create_device_data";
    oi_ssh_device_data* data = NULL;

    data = (oi_ssh_device_data*)malloc(sizeof(oi_ssh_device_data));
    if ( ! data )
    {
        OI_ERROR("%s: Unable to allocate device data structure in memory.",
                 my_name);
        return NULL;
    }

    data->ssh_sock = -1;
    data->ssh_session = NULL;

    data->dev_info.model = NULL;
    data->dev_info.dev_key = NULL;
    data->dev_info.dev_key_len = 0;
    data->dev_info.oi_version = NULL;

	return data;
}

#define safe_free(p) do { if (p) { free((void*)(p)); (p)=NULL;} } while(0)

static void free_device_data(oi_ssh_device_data* data)
{
    if ( ! data )
        return;

    if ( data->ssh_session )
    {
        libssh2_session_disconnect(data->ssh_session, "free_device_data");
        libssh2_session_free(data->ssh_session);
        data->ssh_session = NULL;
    }

    if ( data->ssh_sock )
    {
		closesocket(data->ssh_sock);
        data->ssh_sock = -1;
    }

    safe_free(data->dev_info.model);
    safe_free(data->dev_info.oi_version);

    free(data);
}

static liboi_err_t parse_address(const char* address, struct sockaddr_in* sin)
{
    char* colon = NULL;
    char* port_str = NULL;
    char* ip_addr = NULL;
    int port = 0;

    if ( ! address )
        return EINVAL;

    ip_addr = strdup(address);
    colon = strchr(ip_addr, ':');
    if ( colon )
    {
        *colon = '\0';
        port_str = colon+1;
    }

    port = atoi(port_str);

    sin->sin_family = AF_INET;
    sin->sin_port = htons(port);
    sin->sin_addr.s_addr = inet_addr(ip_addr);

    free(ip_addr);

    return LIBOI_OK;
}

oi_device* oi_ssh_create_device()
{
    oi_device* dev = NULL;

    dev = (oi_device*)malloc(sizeof(oi_device));
    if ( ! dev )
        return NULL;

    memset(dev, 0, sizeof(oi_device));

    dev->dev_data = create_device_data();
    if ( ! dev->dev_data )
    {
        free(dev);
        return NULL;
    }

    return dev;
}

void oi_ssh_free_device(oi_device* device)
{
    if ( ! device )
        return;

    if ( device->dev_data )
        free_device_data(device->dev_data);

    free(device);
}

static liboi_err_t oi_ssh_do_connect(oi_device* device)
{
    const char* my_name = "oi_ssh_connect";
    int ret = 0;
    liboi_err_t oi_err = LIBOI_OK;
    struct sockaddr_in sin;
    oi_ssh_device_data* dev_data;
    const char* dev_key = NULL;

    dev_data = device->dev_data;

    oi_err = parse_address(device->dev_loc.address, &sin);
    if ( LIBOI_OK != oi_err )
    {
        OI_ERROR("%s: Unable to parse address %s.", my_name,
                 device->dev_loc.address);
        goto error_clean;
    }

    dev_data->ssh_sock = socket(AF_INET, SOCK_STREAM, 0);
    if ( dev_data->ssh_sock == -1 )
    {
        OI_ERROR("%s: Unable to create socket.", my_name);
        oi_err = errno;
        goto error_clean;
    }

    dev_data->ssh_session = libssh2_session_init();
    if ( dev_data->ssh_session == NULL )
    {
        OI_ERROR("%s: Unable to create ssh session.", my_name);
        oi_err = errno;
        goto error_clean;
    }

    ret = connect(dev_data->ssh_sock,
                  (struct sockaddr*)(&sin),
                  sizeof(sin));
    if ( ret < 0 )
    {
        OI_ERROR("%s: connect address %s failed: %d, %s.", my_name,
                 device->dev_loc.address, errno,
                 strerror(errno));
        oi_err = errno;
        goto error_clean;
    }

    ret = libssh2_session_startup(dev_data->ssh_session,
                                  dev_data->ssh_sock);
    if ( ret != 0 )
    {
        oi_err = OI_SSH_LASTERROR(my_name, dev_data, "start SSH2 session failed");
        goto error_clean;
    }

    OI_NORMAL("%s: SSH2 session started to OI device.", my_name);

    dev_data->dev_info.dev_key = 
        libssh2_hostkey_hash(dev_data->ssh_session, LIBSSH2_HOSTKEY_HASH_SHA1);
    dev_data->dev_info.dev_key_len = 20;

    ret = libssh2_userauth_password(dev_data->ssh_session, "root", "zxc");
    if ( ret != 0 )
    {
        oi_err = OI_SSH_LASTERROR(my_name, dev_data, "SSH2 authentication failed");
        goto error_clean;
    }

    return LIBOI_OK;

  error_clean:
    if ( dev_data->ssh_session != NULL )
    {
        libssh2_session_free(dev_data->ssh_session);
        dev_data->ssh_session = NULL;
    }
    
    if ( dev_data->ssh_sock != -1 )
    {
        closesocket(dev_data->ssh_sock);
        dev_data->ssh_sock = -1;
    }

    return oi_err;
}

liboi_err_t oi_ssh_connect(oi_device* device)
{
    oi_ssh_device_data* dev_data;

    if ( ! device || ! device->dev_data )
        return EINVAL;

    dev_data = device->dev_data;

    if ( dev_data->ssh_sock != -1 )
    {
        return LIBOI_E_CONNECTED;
    }

    return oi_ssh_do_connect(device);
}

liboi_err_t oi_ssh_disconnect(oi_device* device)
{
    const char* my_name = "oi_ssh_disconnect";
    liboi_err_t oi_err = LIBOI_OK;
    oi_ssh_device_data* dev_data;

    if ( ! device || ! device->dev_data )
        return EINVAL;

    dev_data = device->dev_data;
    
    if ( dev_data->ssh_session != NULL )
    {
        libssh2_session_disconnect(dev_data->ssh_session, "oi_ssh_disconnect");
    }

    if ( dev_data->ssh_sock != -1 )
    {
        closesocket(dev_data->ssh_sock);
    }

    return LIBOI_OK;
}

static int is_dev_connected(oi_ssh_device_data* dev_data)
{
    if (dev_data->ssh_sock != -1)
        return 1;
    else
        return 0;
}

#define BUF_SIZE 4096

static liboi_err_t oi_ssh_write_stdin(LIBSSH2_CHANNEL* channel,
                                      const struct data_buffer* stdin_)
{
    const char* my_name = "oi_ssh_write_stdin";
    ssize_t written = 0;

    assert (channel != NULL);
    assert (stdin_ != NULL);
    assert (stdin_->length > 0);
    assert (stdin_->data != NULL);

    while ( written < stdin_->length )
    {
        ssize_t ioret = libssh2_channel_write(channel,
                                              stdin_->data + written,
                                              stdin_->length - written);
        if ( ioret < 0 )
        {
            OI_ERROR("%s: Write stdin failed: %d. written=%d, towrite=%d.",
                     my_name, ioret, written, stdin_->length);
            return ioret;
        }

        written += ioret;
    }

    return LIBOI_OK;
}

static liboi_err_t oi_ssh_read_output(LIBSSH2_CHANNEL* channel,
                                      int expect_size,
                                      int stream_id,
                                      struct data_buffer* output)
{
    const char* my_name = "oi_ssh_read_output";

    char* buf = (char*)malloc(BUF_SIZE);
    int buf_size = BUF_SIZE;
    int used_size = 0;
    int ioret = 0;
    
    memset(buf, 0, buf_size);

    while ( expect_size == 0 || expect_size > used_size )
    {
        int free_size = buf_size-used_size-1;
        ioret = libssh2_channel_read_ex(channel, stream_id, buf+used_size, free_size);
        if ( ioret < 0 )
        {
            OI_ERROR("%s: Read channel failed: %d. stream_id=%d, free_size=%d",
                     my_name, ioret, stream_id, free_size);
            return ioret;
        }
        else if ( ioret == 0)
        {
            break;
        }
        used_size += ioret;

        if (  buf_size-used_size == 1 )
        {
            int new_buf_size = buf_size + BUF_SIZE;
            buf = (char*)realloc(buf, new_buf_size);
            memset(buf+buf_size, 0, BUF_SIZE);
            buf_size = new_buf_size;
        }
    }

    if ( output )
    {
        output->data = buf;
        output->length = used_size;
    }
    else
    {
        // Ignore the data, should free it here.
        free(buf);
    }

    return LIBOI_OK;
}

liboi_err_t oi_ssh_exec_in_device(oi_device* dev,
                                  const char* command,
                                  struct data_buffer* stdin_,
                                  int* exit_code,
                                  struct data_buffer* stdout_,
                                  struct data_buffer* stderr_)
{
    const char* my_name = "oi_ssh_execute_in_device";
    oi_ssh_device_data* dev_data;
    liboi_err_t oi_err = LIBOI_OK;
    int err_no = 0;
    LIBSSH2_CHANNEL* channel = NULL;

    if ( ! dev || ! dev->dev_data || ! command )
        return EINVAL;

    dev_data = dev->dev_data;

    if ( ! is_dev_connected(dev_data) )
        return LIBOI_E_NOT_CONNECTED;

    OI_NORMAL("%s: create new channel.", my_name);

    channel = libssh2_channel_open_session(dev_data->ssh_session);
    if ( channel == NULL )
        return OI_SSH_LASTERROR(my_name, dev_data, "create channel failed");

    err_no = libssh2_channel_exec(channel, command);
    if ( err_no != 0 )
    {
        libssh2_channel_free(channel);
        OI_ERROR("%s: execute command failed: %d, %s.", my_name,
                 err_no, command);
        return err_no;
    }

    if ( stdin_ && stdin_->length && stdin_->data )
    {
        err_no = oi_ssh_write_stdin(channel, stdin_);
        if ( err_no != LIBOI_OK )
        {
            libssh2_channel_free(channel);
            OI_ERROR("%s: Write stdin failed: %d, %s, stdin=[%d]%s.", my_name,
                     err_no, command, stdin_->length, stdin_->data);
            return err_no;
        }
    }

    err_no = oi_ssh_read_output(channel, 0, 0, stdout_);
    if ( err_no != LIBOI_OK )
    {
        libssh2_channel_free(channel);
        return err_no;
    }

    err_no = oi_ssh_read_output(channel, 0, SSH_EXTENDED_DATA_STDERR, stderr_);
    if ( err_no != LIBOI_OK )
    {
        libssh2_channel_free(channel);
        oi_free_data_buffer(stdout_);
        return err_no;
    }

    err_no = libssh2_channel_close(channel);
    if ( err_no < 0 )
    {
        OI_WARN("%s: Close channel failed: %d.", my_name, err_no);
    }
    
    if ( exit_code )
        *exit_code = libssh2_channel_get_exit_status(channel);

    libssh2_channel_free(channel);

    return LIBOI_OK;
}

static char* oi_ssh_substring(const char* start, const char* end)
{
    if ( end == NULL )
    {
        return strdup(start);
    }
    else
    {    
        int len = end - start;
        char* buf = (char*)malloc(len + 1);
        strncpy(buf, start, len);
        buf[len] = '\0';
        return buf;
    }
}

static liboi_err_t oi_ssh_parse_version(oi_device_info* dev_info, const char* version)
{
#define LEN_OF_OPENINKPOT (sizeof("OpenInkpot")-1)

    const char* my_name = "oi_ssh_parse_version";
    const char* version_start = NULL;
    const char* version_end = NULL;
    const char* model_start = NULL;
    const char* model_end = NULL;

    dev_info->model = NULL;
    dev_info->oi_version = NULL;

    if ( strncmp(version, "OpenInkpot", LEN_OF_OPENINKPOT) != 0 ||
         version[LEN_OF_OPENINKPOT] != ' ')
    {
        OI_DEBUG("%s: Version string not start with OpenInkpot.", my_name);
        return LIBOI_E_NOT_OPENINKPOT;
    }
    
    version_start = version+LEN_OF_OPENINKPOT+1;
    version_end = strchr(version_start, ' ');
    dev_info->oi_version = oi_ssh_substring(version_start, version_end);
    OI_DEBUG("%s: oi_version: %d-%d, %s", my_name,
             version_start-version,
             (version_end ? version_end-version : -1),
             dev_info->oi_version);
    if ( version_end == NULL )
        return LIBOI_OK;

    model_start = version_end + 1;
    if ( *model_start == '(' )
    {
        model_start++;
        model_end = strchr(model_start, ')');
    }

    dev_info->model = oi_ssh_substring(model_start, model_end);
    OI_DEBUG("%s: oi_model: %d-%d, %s", my_name,
             model_start-version,
             (model_end ? model_end-version : -1),
             dev_info->model);

    return LIBOI_OK;
}

liboi_err_t oi_ssh_get_info(oi_device* dev, oi_device_info* dev_info)
{
    const char* my_name = "oi_ssh_get_info";
    liboi_err_t oi_err = LIBOI_OK;
    oi_ssh_device_data* dev_data = NULL;
    struct data_buffer oi_version = {0};
    struct data_buffer oi_ver_err = {0};
    int exit_code = 0;

    if ( ! dev || ! dev->dev_data || ! dev_info )
        return EINVAL;

    dev_data = dev->dev_data;

    if ( ! is_dev_connected(dev_data) )
        return LIBOI_E_NOT_CONNECTED;

    oi_err = oi_ssh_exec_in_device(dev,
                                   "cat /etc/openinkpot-version",
                                   NULL,
                                   &exit_code,
                                   &oi_version,
                                   &oi_ver_err);
    if ( oi_err != LIBOI_OK || exit_code != 0 )
    {
        OI_ERROR("%s: Unable to get OI version file /etc/openinkpot-version:"
                 " [%d][%d] %s.",
                 my_name, oi_err, exit_code, oi_ver_err.data);
        oi_free_data_buffer(&oi_version);
        oi_free_data_buffer(&oi_ver_err);
        return oi_err;
    }

    oi_free_data_buffer(&oi_ver_err);
    OI_NORMAL("%s: /etc/openinkpot-version: %s", my_name, oi_version.data);

    dev_info->dev_key = dev_data->dev_info.dev_key;
    dev_info->dev_key_len = dev_data->dev_info.dev_key_len;

    oi_err = oi_ssh_parse_version(dev_info, oi_version.data);
    
    oi_free_data_buffer(&oi_version);
    return oi_err;
}

liboi_err_t oi_ssh_get_file(oi_device* device,
                            const char* device_file_path,
                            const char* host_file_path)
{
    const char* my_name = "oi_ssh_get_file";
    liboi_err_t oi_err = LIBOI_OK;
    int io_ret = 0;
    oi_ssh_device_data* dev_data = NULL;
    struct data_buffer file_data = {0};
    LIBSSH2_CHANNEL* channel = NULL;
    struct stat file_stat = {0};
    FILE* output = NULL;

    if ( ! device || ! device->dev_data ||
         ! device_file_path || ! host_file_path )
        return EINVAL;

    dev_data = device->dev_data;
    if ( ! is_dev_connected(dev_data) )
        return LIBOI_E_NOT_CONNECTED;

    channel = libssh2_scp_recv(dev_data->ssh_session,
                               device_file_path,
                               &file_stat);
    if ( channel == NULL )
        return OI_SSH_LASTERROR(my_name, dev_data, "get file failed");

    output = fopen(host_file_path, "wb");
    if ( output == NULL )
    {
        OI_ERROR("%s: Unable to write to output file: %d: %s", my_name,
                 errno, strerror(errno));
        libssh2_channel_free(channel);
        return errno;
    }

    oi_err = oi_ssh_read_output(channel, file_stat.st_size, 0, &file_data);
    if ( oi_err != LIBOI_OK )
    {
        libssh2_channel_free(channel);
        fclose(output);
        return oi_err;
    }

    if ( file_stat.st_size != file_data.length )
    {
        OI_ERROR("%s: Data read from scp length does not match:"
                 " should be %d bytes, but read %d bytes.",
                 my_name, file_stat.st_size, file_data.length);
        libssh2_channel_free(channel);
        oi_free_data_buffer(&file_data);
        fclose(output);
        return oi_err;
    }

    libssh2_channel_free(channel);
    channel = NULL;

    io_ret = fwrite(file_data.data,
                    file_data.length,
                    1, output);
    if ( io_ret != 1 )
    {
        OI_ERROR("%s: While write data with length %d to file %s,"
                 " got an write error:%d:%s.",
                 my_name, file_data.length, host_file_path,
                 errno, strerror(errno));
        oi_free_data_buffer(&file_data);
        fclose(output);
        return errno;
    }

    oi_free_data_buffer(&file_data);
    fclose(output);

    OI_NORMAL("%s: File on device %s was successfully copied to %s.",
              my_name, device_file_path, host_file_path);

    return LIBOI_OK;
}

liboi_err_t oi_ssh_put_file(oi_device* device,
                            const char* host_file_path,
                            const char* device_file_path)
{
    const char* my_name = "oi_ssh_put_file";
    liboi_err_t oi_err = LIBOI_OK;
    oi_ssh_device_data* dev_data = NULL;
    LIBSSH2_CHANNEL* channel = NULL;
    FILE* input = NULL;
    struct data_buffer file_data = {0};
    char data_buffer[4096];
    int io_ret = 0;
    struct stat file_stat = {0};

    if ( ! device || ! device->dev_data ||
         ! device_file_path || ! host_file_path )
        return EINVAL;

    dev_data = device->dev_data;
    if ( ! is_dev_connected(dev_data) )
        return LIBOI_E_NOT_CONNECTED;

    io_ret = stat(host_file_path, &file_stat);
    if ( io_ret < 0 )
    {
        OI_ERROR("%s: Unable to get file stat %s: %d: %s", my_name,
                 host_file_path, errno, strerror(errno));
        return errno;
    }

    channel = libssh2_scp_send_ex(dev_data->ssh_session,
                                  device_file_path,
                                  file_stat.st_mode & 0xFF,
                                  file_stat.st_size,
                                  file_stat.st_mtime,
                                  file_stat.st_atime);
    if ( channel == NULL )
        return OI_SSH_LASTERROR(my_name, dev_data, "put file failed");

    input = fopen(host_file_path, "rb");
    if ( input == NULL )
    {
        OI_ERROR("%s: Unable to read %s: %d: %s", my_name,
                 host_file_path, errno, strerror(errno));
        libssh2_channel_free(channel);
        return errno;
    }

    file_data.data = data_buffer;
    file_data.length = sizeof(data_buffer);
    while ( ! feof(input) )
    {
        io_ret = fread(file_data.data,
                       1,
                       file_data.length,
                       input);
        if ( io_ret <= 0 )
        {
            OI_ERROR("%s: While reading data with length %d from file %s,"
                     " got an write error:%d:%s.",
                     my_name, file_data.length, host_file_path,
                     errno, strerror(errno));
            fclose(input);
            libssh2_channel_free(channel);
            return errno;
        }

        oi_ssh_write_stdin(channel, &file_data);
    }

    fclose(input);
    libssh2_channel_free(channel);

    OI_NORMAL("%s: File %s was successfully copied to %s.",
              my_name, host_file_path, device_file_path);

    return LIBOI_OK;
}

